#  -*- coding: utf-8 -*-
import pandas as pd
from pymongo import UpdateOne,ASCENDING, DESCENDING
from factor.base_factor import BaseFactor
from data.finance_report_crawler import FinanceReportCrawler
from data.data_module import DataModule
from util.stock_util import get_all_codes,get_all_indexes_date,calc_negative_diff_dates,multi_computer,get_all_codes_trading_date,get_trading_dates
from util.database import DB_CONN
import time
from datetime import datetime, timedelta

"""
实现涨幅中位数的因子计算和保存
"""


class MedianFactor(BaseFactor):
    def __init__(self):
        BaseFactor.__init__(self, name='median')

    def compute_median_date(self,date,is_index):
        start_time = time.time()
        update_requests = []

        factor_cursor = DB_CONN['change_rate'].find(
            {'date':date, "index": is_index},
            projection={'_id': False})

        factor_df = pd.DataFrame(
            [x for x in factor_cursor])

        hs_median_close_change_rate = factor_df.median(skipna  =  True)['close_change_rate']

        update_requests.append(
            UpdateOne(
                {'date': date},
                {'$set': {'hs_median_close_change_rate': hs_median_close_change_rate}},
                upsert=True))


        if len(update_requests) > 0:
            update_result = self.collection.bulk_write(update_requests, ordered=False)
            end_time = time.time()
            print('填充Median因子，日期：%s，插入：%4d条，更新：%4d条,耗时：%.3f 秒' %
                  (date, update_result.upserted_count, update_result.modified_count, (end_time - start_time)),
                  flush=True)

    def compute(self, begin_date, end_date):
        """
        计算指定时间段内所有股票的该因子的值(中位数)，并保存到数据库中
        :param begin_date:  开始时间
        :param end_date: 结束时间
        """
        self.collection.create_index([('date', 1)])

        #获取所有股票
        dates = get_trading_dates(begin_date, end_date)
        args = (False,)
        multi_computer(compute_dates, dates, args)


def compute_dates(dates,is_index):
    for date in dates:
        MedianFactor().compute_median_date(date, is_index)



if __name__ == '__main__':
    # 执行因子的提取任务
    #hfq =HfqMAFactor()
    pd.set_option('display.width',500)
    pd.set_option('display.max_columns', 500)
    pd.set_option('display.max_colwidth', 500)
    MedianFactor().compute("2010-01-01", '2020-12-11')
